package wikiextract.nlp.trainingset.x;

import static wikiextract.nlp.trainingset.x.Settings.NUM_CROSS_LEXICONS;
import static wikiextract.nlp.trainingset.x.Settings.WEB_LEXICONS;
import static wikiextract.nlp.trainingset.x.RunSecondPass.secondFeatureIds;
import static wikiextract.nlp.trainingset.x.RunSecondPass.featureIds;
import static wikiextract.nlp.trainingset.x.RunSecondPass.getSecondFeatureId;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import wikiextract.data.util.StringUtil;

public class SecondPassLexicons {

	static class LexiconFeature {
		String featureName;
		int featureId;
		Lexicons lexicons;
		LexiconFeature(String featureName) {
			this.featureName = featureName;
			this.featureId = getSecondFeatureId(featureName, secondFeatureIds, featureIds);
			lexicons = new Lexicons();
		}
	}
	
	static class Lexicons {
		Lexicon train;
		//Lexicon test;
		List<Lexicon> chunkTrain;
		List<Lexicon> chunkTest;
		
		Lexicons() {
			train = new Lexicon();
			//test = new Lexicon();
			chunkTrain = new ArrayList<Lexicon>();
			chunkTest = new ArrayList<Lexicon>();
			for (int c = 0; c < NUM_CROSS_LEXICONS; c++) {
				chunkTrain.add(new Lexicon());
				chunkTest.add(new Lexicon());
			}
		}
	}
	
	static class Lexicon {
		Map<String, Integer> lex = new HashMap<String,Integer>();
		Map<String, List<String[]>> index = null;
	}

	static void seeds(List<List<FeaturesRecord>> train, 
			List<LexiconFeature> lexicons, String featureName, int label, int word) {
		LexiconFeature lf = new LexiconFeature(featureName);		
		Lexicons lex = lf.lexicons;
			
		// compute the entirety of all values outside the current cv set
		for (List<FeaturesRecord> articleExamples : train)
			add(articleExamples, lex.train.lex, label, word);

		//for (List<Output> articleExamples : test)
		//	add(articleExamples, lex.test.lex);
				
		// break up trainingExamples into max NUM_CROSS_LISTS chunks
		// assign articles to chunks
		for (int k = 0; k < train.size(); k++) {
			int c = k % NUM_CROSS_LEXICONS;

			// add seeds to all other chunks, current one only for testing
			for (int d = 0; d < Math.min(NUM_CROSS_LEXICONS, train.size()); d++) {					 
				if (c != d) {
					add(train.get(k), lex.chunkTrain.get(d).lex, label, word);
				}
				//else {
				//	add(train.get(k), lex.chunkTest.get(d).lex);
				//}
			}
		}
		lexicons.add(lf);
	}

	public static final byte CMD_EXPAND = 0;
	public static final byte CMD_QUIT = 1;
	
    final static String server = "pardosa03";
    final static int port = 9999;
    //final static byte[] serverIp = {128, 208, 2, 134};
    
    static InetAddress serverA = null;
    static {
    	try {
    		serverA = InetAddress.getByName(server);
    	} catch (UnknownHostException e) { e.printStackTrace(); }
    }
    
	static void weblists(LexiconFeature seeds, List<LexiconFeature> lexicons,
			String featureNamePrefix) {
		List<LexiconFeature> webLexicons = new ArrayList<LexiconFeature>(WEB_LEXICONS);
		for (int lex = 0; lex < WEB_LEXICONS; lex++)
			webLexicons.add(new LexiconFeature(featureNamePrefix + lex));

        DataInputStream is = null;
        DataOutputStream os = null;
        Socket sock = null;
        try {
            sock = new Socket(serverA, port);
            is = new DataInputStream(sock.getInputStream());
            os = new DataOutputStream(sock.getOutputStream());

            Lexicons seedL = seeds.lexicons;
				
			List<Lexicon> tgts = new ArrayList<Lexicon>();
			for (int l=0; l < WEB_LEXICONS; l++) 
				tgts.add(webLexicons.get(l).lexicons.train);	
			webHelper(seedL.train, tgts, is, os);

//			tgts.clear();
//			for (int l=0; l < WEB_LEXICONS; l++) 
//				tgts.add(webLexicons.get(l).lexicons.test);
//			webHelper(seedL.test, tgts, is, os);

			for (int c = 0; c < NUM_CROSS_LEXICONS; c++) {
				tgts.clear();
				for (int l=0; l < WEB_LEXICONS; l++) 
					tgts.add(webLexicons.get(l).lexicons.chunkTrain.get(c));
				webHelper(seedL.chunkTrain.get(c), tgts, is, os);

				tgts.clear();
				for (int l=0; l < WEB_LEXICONS; l++) 
					tgts.add(webLexicons.get(l).lexicons.chunkTest.get(c));
				webHelper(seedL.chunkTest.get(c), tgts, is, os);				
			}
			lexicons.addAll(webLexicons);
		} catch (Throwable e) {
		     System.out.println("Error " + e.getMessage());
		     e.printStackTrace();
		} finally {
			try { is.close(); } catch (Exception e) { e.printStackTrace(); }
			try { os.close(); } catch (Exception e) { e.printStackTrace(); }
			try { sock.close(); } catch (Exception e) { e.printStackTrace(); }
		}
	}
			
	private static void webHelper(Lexicon seeds, List<Lexicon> targets,
			DataInputStream is, DataOutputStream os) throws Exception {
		
		// expand train
	    os.write(CMD_EXPAND);
	    os.writeInt(seeds.lex.size());
	    for (Map.Entry<String,Integer> en : seeds.lex.entrySet()) {
	    	os.writeInt(en.getValue()); // freq
	    	os.writeUTF(en.getKey()); // word
	    }
    
        for (int lex=0; lex < WEB_LEXICONS; lex++) {        	
        	Map<String,Integer> target = targets.get(lex).lex;
        	int size = is.readInt();
        	for (int i =0; i < size; i++) {
        		int freq = is.readInt();
        		String t = is.readUTF();
        		target.put(t, freq);
        		//System.out.println(t);
        	}
        }
	}
	
	static void createIndexes(List<LexiconFeature> lexicons) 
		throws IOException {
		for (LexiconFeature lf : lexicons) {
			Lexicons ls = lf.lexicons;
			ls.train.index = createIndex(ls.train.lex);
			//	cvs.test.index = createIndex(cvs.test.lex);				
			for (Lexicon l : ls.chunkTrain)
				l.index = createIndex(l.lex);
			for (Lexicon l : ls.chunkTest)
				l.index = createIndex(l.lex);
		}
	}

	static void lexiconFeatures(FeaturesRecord o, boolean test, Map<String,List<String[]>> lexicon,
			List<Integer>[] fts, int featureId) {
		for (int i = 0; i < o.tokens.length; i++) {
			String tok = o.tokens[i].toLowerCase();

			
			// CHECK: why is below the same in both branches?
			if (!test) {
				// training data, consider chunks
				List<String[]> matchingSegmentContained = lexicon.get(tok);
				if (matchingSegmentContained != null) {
					for (String[] segment : matchingSegmentContained) {					
						if (segmentMatch(o.tokens, i, segment)) {	
							fts[i].add(featureId);
							break;
						}
					}
				}
			} else {
				// test data
				List<String[]> matchingSegmentContained = lexicon.get(tok);
				if (matchingSegmentContained != null) {
					for (String[] segment : matchingSegmentContained) {							
						if (segmentMatch(o.tokens, i, segment)) {
							fts[i].add(featureId);
							break;
						}
					}
				}
			}
		}
	}

	private static boolean segmentMatch(String[] tokens, int pos, String[] segment) {
		for (int k=Math.max(0, pos-segment.length); k <= Math.min(pos, tokens.length-segment.length); k++) {
			boolean match = true;
			for (int p = 0; p < segment.length; p++)
				if (!tokens[k+p].toLowerCase().equals(segment[p])) {
					match = false;
					break;
				}
			if (match) return true;
		}
		return false;
	}
	
	// word: 0 = all, 1 = only first word of segment
	static void add(List<FeaturesRecord> article, Map<String,Integer> seeds,
			int label, int word) {
		for (FeaturesRecord ex : article) {
			for (int m = 0; m < ex.matchPos.length; m++) {
				if (ex.matchLabels[m] != label) continue;
				byte[] t = ex.matchPos[m];
				StringBuilder sb = new StringBuilder();
				for (int j = t[0]; j < t[1]; j++) {					
					if (j > t[0])
						sb.append(" ");
					sb.append(ex.tokens[j].toLowerCase());
					if (word == 1) break; // only first word
				}
				increment(seeds, sb.toString());
			}
		}	
	}
	
	private static void increment(Map<String, Integer> map, String key) {
		Integer i = map.get(key);
		if (i == null)
			map.put(key, 1);
		else
			map.put(key, i + 1);
	}
		
	static HashMap<String, List<String[]>> createIndex(Map<String,Integer> freqs) throws IOException {
		HashMap<String, List<String[]>> index = new HashMap<String, List<String[]>>();
		for (Map.Entry<String,Integer> e : freqs.entrySet()) {
			String[] toks = StringUtil.split(e.getKey(), ' ', 10);
			for (String tok : toks) {
				List<String[]> s = index.get(tok);
				if (s == null) {
					s = new ArrayList<String[]>(3);
					index.put(tok, s);
				}
				s.add(toks);
				// there might be duplicates ...
			}					
		}
		return index;
	}

}
